from django.contrib.auth import get_user_model
from user.models import user as user_model
from functools import wraps
from .custom_response import custom_response

UserModel = get_user_model()


def login_required(func):
    @wraps(func)
    def wrapped_login_required(request):
        user_token = request.headers.get('Oauth-Token')
        user_info = user_model.objects.filter(token__token=user_token).values('id', 'account', 'nickname',
                                                                              'avatar').first()
        if not user_info:
            resp = custom_response()
            resp.set_code(1005)
            resp.set_message('未查询到用户或登录超时！')
            return resp.response
        return func(request)

    return wrapped_login_required
